-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2025 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module Test.Auth where

import API.Brig
import API.BrigInternal
import API.Common
import API.GalleyInternal
import qualified API.Nginz as Nginz
import qualified Data.ByteString.Char8 as BSChar8
import SetupHelpers
import Testlib.Prelude
import Text.Read
import UnliftIO.Async
import UnliftIO.Concurrent

-- Happy flow: login yields a valid zauth token.
--
-- See also: 'testBearerToken'
testBearerToken2 :: (HasCallStack) => App ()
testBearerToken2 = do
  alice <- randomUser OwnDomain def
  email <- asString $ alice %. "email"
  loginResp <- login alice email defPassword >>= getJSON 200
  token <- asString $ loginResp %. "access_token"

  req <-
    rawBaseRequest alice Nginz Versioned "/self"
      <&> addHeader "Authorization" ("Bearer " <> token)
  submit "GET" req `bindResponse` \resp -> do
    resp.status `shouldMatchInt` 200
    resp.json %. "email" `shouldMatch` email

-- Happy flow (zauth token encoded in AWS4_HMAC_SHA256)
--
-- See also: 'testAWS4_HMAC_SHA256_token'
testAWS4_HMAC_SHA256_token2 :: (HasCallStack) => App ()
testAWS4_HMAC_SHA256_token2 = do
  alice <- randomUser OwnDomain def
  email <- asString $ alice %. "email"
  loginResp <- login alice email defPassword >>= getJSON 200
  token <- asString $ loginResp %. "access_token"

  let testCases =
        [ (True, "AWS4-HMAC-SHA256 Credential=" <> token <> ", foo=bar"),
          (True, "AWS4-HMAC-SHA256 Credential=" <> token),
          (True, "AWS4-HMAC-SHA256 foo=bar, Credential=" <> token),
          (True, "AWS4-HMAC-SHA256 foo=bar, Credential=" <> token <> ", baz=qux"),
          (True, "AWS4-HMAC-SHA256 foo=bar,Credential=" <> token <> ",baz=qux"),
          (False, "AWS4-HMAC-SHA256 Credential=badtoken")
        ]

  for_ testCases $ \(good, header) -> do
    req <-
      rawBaseRequest alice Nginz Versioned "/self"
        <&> addHeader "Authorization" header
    submit "GET" req `bindResponse` \resp -> do
      if good
        then do
          resp.status `shouldMatchInt` 200
          resp.json %. "email" `shouldMatch` email
        else do
          resp.status `shouldMatchInt` 401

-- The testLimitRetries test conforms to the following testing standards:
-- @SF.Channel @TSFI.RESTfulAPI @TSFI.NTP @S2
--
-- The following test tests the login retries. It checks that a user can make
-- only a prespecified number of attempts to log in with an invalid password,
-- after which the user is unable to try again for a configured amount of time.
-- After the configured amount of time has passed, the test asserts the user can
-- successfully log in again. Furthermore, the test asserts that another
-- unrelated user can successfully log-in in parallel to the failed attempts of
-- the aforementioned user.
testLimitRetries :: (HasCallStack) => App ()
testLimitRetries = do
  let retryLimit = 5
      timeout = 5
  withModifiedBackend
    def
      { brigCfg =
          -- Set a small timeout to make this test fast
          setField @_ @Int "optSettings.setLimitFailedLogins.timeout" timeout
            >=> setField @_ @Int "optSettings.setLimitFailedLogins.retryLimit" retryLimit
            -- Disable password hashing rate limiting, so we can login many times without making this test slow
            >=> setField @_ @Int "optSettings.setPasswordHashingRateLimit.userLimit.inverseRate" 0
      }
    $ \domain -> do
      alice <- randomUser domain def
      aliceEmail <- asString $ alice %. "email"

      bob <- randomUser domain def
      bobEmail <- asString $ bob %. "email"

      -- Alice tries to login 5 times with wrong password
      forM_ [1 .. retryLimit] $ \_ ->
        login domain aliceEmail "wrong-password" `bindResponse` \resp -> do
          resp.status `shouldMatchInt` 403
          resp.json %. "label" `shouldMatch` "invalid-credentials"

      -- Now alice cannot login even with correct password
      retryAfter <-
        login domain aliceEmail defPassword `bindResponse` \resp -> do
          resp.status `shouldMatchInt` 403
          resp.json %. "label" `shouldMatch` "client-error"
          let Just retryAfter = readMaybe . BSChar8.unpack =<< lookup (fromString "Retry-After") resp.headers
          (retryAfter <= timeout) `shouldMatch` True
          pure retryAfter

      -- Bob can still login
      login domain bobEmail defPassword `bindResponse` \resp -> do
        resp.status `shouldMatchInt` 200

      -- Waiting 2s less than retryAfter should still cause a failure
      threadDelay ((retryAfter - 2) * 1_000_000)
      login domain aliceEmail defPassword `bindResponse` \resp -> do
        resp.status `shouldMatchInt` 403
        resp.json %. "label" `shouldMatch` "client-error"
        let Just retryAfter2 = readMaybe . BSChar8.unpack =<< lookup (fromString "Retry-After") resp.headers
        -- This should be about 2 seconds or slightly less because we didn't
        -- wait long enough. This also asserts that the throttling doesn't get
        -- reset by making another call
        (retryAfter2 <= (2 :: Int)) `shouldMatch` True

      -- Waiting 2 more seconds should make the login succeed
      threadDelay 2_000_000
      login domain aliceEmail defPassword `bindResponse` \resp -> do
        resp.status `shouldMatchInt` 200

-- @END

-- The testTooManyCookies test conforms to the following testing standards:
-- @SF.Provisioning @TSFI.RESTfulAPI @S2
--
-- The test asserts that there is an upper limit for the number of user cookies
-- per cookie type. It does that by concurrently attempting to create more
-- persistent and session cookies than the configured maximum.
-- Creation of new cookies beyond the limit causes deletion of the
-- oldest cookies.
testTooManyCookies :: (HasCallStack) => App ()
testTooManyCookies = do
  let cookieLimit = 5
  withModifiedBackend
    def
      { brigCfg =
          -- Disable password hashing rate limiting, so we can login many times without making this test slow
          setField @_ @Int "optSettings.setPasswordHashingRateLimit.userLimit.inverseRate" 0
            -- Disable cookie throttling so this test is not slow
            >=> setField @_ @Int "optSettings.setUserCookieThrottle.retryAfter" 0
            >=> setField @_ @Int "optSettings.setUserCookieLimit" cookieLimit
      }
    $ \domain -> do
      alice <- randomUser domain def
      aliceEmail <- asString $ alice %. "email"

      let testCookieLimit label = do
            let loginFn = if label == "persistent" then login else loginWithSessionCookie
            (deletedCookie1 : deletedCookie2 : validCookies) <-
              replicateM (cookieLimit + 2)
                $ do
                  -- This threadDelay is required to get around problems caused
                  -- by: https://wearezeta.atlassian.net/browse/WPB-15446
                  threadDelay 1_000_000
                  loginFn domain aliceEmail defPassword
                    `bindResponse` \resp -> do
                      resp.status `shouldMatchInt` 200
                      pure . fromJust $ getCookie "zuid" resp
            addFailureContext ("deletedCookie1: " <> deletedCookie1 <> "\ndeletedCookie2: " <> deletedCookie2 <> "\nvalidCookies:\n" <> unlines validCookies) $ do
              forM_ [deletedCookie1, deletedCookie2] $ \deletedCookie -> do
                renewToken alice deletedCookie `bindResponse` \resp ->
                  resp.status `shouldMatchInt` 403
              forM_ validCookies $ \validCookie ->
                renewToken alice validCookie `bindResponse` \resp ->
                  resp.status `shouldMatchInt` 200
      concurrently_ (testCookieLimit "persistent") (testCookieLimit "session")

-- @END

-- The testInvalidCookie test conforms to the following testing standards:
-- @SF.Provisioning @TSFI.RESTfulAPI @TSFI.NTP @S2
--
-- Test that invalid and expired tokens do not work.
testInvalidCookie :: (HasCallStack) => App ()
testInvalidCookie = do
  let cookieTimeout = 2
  withModifiedBackend
    def
      { brigCfg =
          setField @_ @Int "zauth.authSettings.userTokenTimeout" cookieTimeout
            >=> setField @_ @Int "zauth.authSettings.legalHoldUserTokenTimeout" cookieTimeout
      }
    $ \domain -> do
      Nginz.access domain "zuid=xxx" `bindResponse` \resp -> do
        resp.status `shouldMatchInt` 403
        resp.json %. "label" `shouldMatch` "client-error"
        msg <- asString $ resp.json %. "message"
        msg `shouldContain` "Invalid zauth token"

      (owner, tid, [alice]) <- createTeam domain 2
      aliceEmail <- asString $ alice %. "email"
      aliceId <- asString $ alice %. "qualified_id.id"
      userCookie <-
        login domain aliceEmail defPassword `bindResponse` \resp -> do
          resp.status `shouldMatchInt` 200
          pure . fromJust $ getCookie "zuid" resp

      legalholdWhitelistTeam tid owner >>= assertSuccess

      lhCookie <-
        legalholdLogin domain aliceId defPassword `bindResponse` \resp -> do
          resp.status `shouldMatchInt` 200
          pure . fromJust $ getCookie "zuid" resp

      -- Wait for both cookies to expire
      -- In order to reduce flakiness, the delay has been increased by 2 seconds
      -- after an investigation revealed that check for cookie expiration is somewhat lenient due to
      -- the way the time is calculated in the backend.
      -- See the interpreter of Now which is implemented using `Control.AutoUpdate` which defaults to an update frequency of 1 sec.
      -- Furthermore the timestamp is then again rounded down to the nearest second before it is compared to the cookie expiration time.
      threadDelay $ (cookieTimeout + 2) * 1_000_000
      -- Assert that the cookies are considered expired
      for_ [userCookie, lhCookie] $ \cookie ->
        Nginz.access domain ("zuid=" <> cookie) `bindResponse` \resp -> do
          resp.status `shouldMatchInt` 403
          resp.json %. "label" `shouldMatch` "invalid-credentials"
          resp.json %. "message" `shouldMatch` "Zauth token expired"

-- @END
